'''
Created on Apr 6, 2012

@author: Rafael
'''
import os
import time
from nipy import load_image
from scipy.cluster.vq import kmeans2
import numpy as np
import scipy as sp
from scipy.stats import pearsonr as pearson

bd = './Inputs'
os.chdir(bd)
infile='Sub_1_proc.nii'
inmask='Sub_1_GM.nii.gz'

in_nii = load_image(infile)
in_mask = load_image(inmask)

data=in_nii.get_data()
mask=in_mask.get_data()

shape=data.shape

log=open('k_prep.log', 'w')


print "Data dimensions: " , data.shape
log.write("Data dimensions: " , data.shape)
print "Mask dimensions: ", mask.shape
log.write("Mask dimensions: ", mask.shape)


voxels=shape[0]*shape[1]*shape[2]
print 'original voxels: ',  voxels
log.write('original voxels: ',  voxels)


TS_Vector=np.zeros( (voxels, shape[3]) )
MS_Vector=np.zeros( ( voxels) )

i = 0
while i < voxels-1:
    for z in range(shape[2]):
        for y in range(shape[1]):
            for x in range(shape[0]):
                TS_Vector[i, : ]=data[x,y,z,:]
                MS_Vector[i]=mask[x,y,z]
                i+=1

GM = sum(1 for i in MS_Vector if i)

print "number of GM voxels: " , GM
log.write("number of GM voxels: " , GM)

GM_Vector= np.zeros( ( GM, shape[3] ))

j=0
for i in range(MS_Vector.shape[0]):
    if MS_Vector[i] != 0:
        GM_Vector[j]=TS_Vector[i, :]
        j+=1


print "Vector Shape" ,GM_Vector.shape
log.write("Vector Shape" ,GM_Vector.shape)

print "Started correlations at: ",  time.ctime()
log.write("Started correlations at: ",  time.ctime())

#Corr_Matrix=sp.sparse.lil_matrix((GM, GM))      
Corr_Matrix=np.zeros( (GM, GM) )

TotCorr=(GM*GM)

k=-1
counter=0
total=0
for i in range(Corr_Matrix.shape[0]):
    k+=1
    for j in range(k,Corr_Matrix.shape[1]):
        c=pearson(GM_Vector[i,:],GM_Vector[j,:])[0]
        Corr_Matrix[i, j] = c
        Corr_Matrix[j, i] = c
        counter+=1
        total+=1
        if counter == 10000:
            print "corr: ", i, j
            print "Percent ",TotCorr
            counter=0

print "Corr matrix shape" , Corr_Matrix.shape
log.write("Corr matrix shape" , Corr_Matrix.shape)
print Corr_Matrix[1309, 1309]
print Corr_Matrix[1308, 1178]

print Corr_Matrix.shape

np.save('correlation_matrix', Corr_Matrix)

print "End correlations at: " , time.ctime()
log.write("End correlations at: " , time.ctime())
log.close()


